#!/usr/bin/env python2
# coding: utf-8

import dbus
import sys
import re

from os import path as Path

from pynotify import URGENCY_CRITICAL
from pynotify import URGENCY_LOW
from pynotify import URGENCY_NORMAL

DBUS_PROPS_INTERFACE = 'org.freedesktop.DBus.Properties'
UDISKS_INTERFACE = 'org.freedesktop.UDisks'
UDISKS_DEVICE_INTERFACE = 'org.freedesktop.UDisks.Device'

UDISKS_OBJECT = 'org.freedesktop.UDisks'
UDISKS_OBJECT_PATH = '/org/freedesktop/UDisks'

ICON_SUCCESS = '/usr/share/dmenu_udisks/USB-HD.png'
ICON_ERROR = '/usr/share/dmenu_udisks/USB-HD_err.png'

class Udisks:
    def __init__(self):
        self.bus = dbus.SystemBus()
        self.ud_manager_obj = self.bus.get_object(UDISKS_OBJECT, UDISKS_OBJECT_PATH)
        self.ud_manager = dbus.Interface(self.ud_manager_obj, UDISKS_OBJECT)

        #device_obj = self.bus.get_object(UDISKS_OBJECT, dev)
        #self.device_props = dbus.Interface(self.device_obj, dbus.PROPERTIES_IFACE)
        self.Properties = self._get_all_properties()

    def get_udisk_devs(self):
        udisks = self.bus.get_object(UDISKS_OBJECT, UDISKS_OBJECT_PATH)
        return udisks.EnumerateDevices(dbus_interface=UDISKS_INTERFACE)


    def _get_all_properties(self):
        'Retunt Dictionay of dev with props'
        devlist = self.get_udisk_devs()
        
        d = dbus.Dictionary()
        for i in devlist:
            dev_obj = self.bus.get_object(UDISKS_OBJECT, i)
            dev = dbus.Interface(dev_obj, DBUS_PROPS_INTERFACE)
            d[i] = dev.GetAll('')
        return d

    # --- PARAMETERS ---
    def _get_property(self, dev, property):
        return self.Properties[dev].get(property, None) if dev in self.Properties.keys() else None

    def is_mounted(self, device):
        return bool(self._get_property(device, 'DeviceIsMounted'))

    def mount_paths(self, device):
        raw_paths = self._get_property(device, 'DeviceMountPaths')
        return [Path.normpath(path) for path in raw_paths]

    def is_filesystem(self, device):
        return bool(self._get_property(device, 'IdUsage') == 'filesystem')

    def has_media(self, device):
        return bool(self._get_property(device, 'DeviceIsMediaAvailable'))

    def id_type(self, device):
        return str(bool(self._get_property(device, 'IdType')))

    def devlabel(self, device):
        return str(self._get_property(device, 'IdLabel'))

    def devname(self, device):
        return str(self._get_property(device, 'DeviceFile'))

    def deviface(self, device):
        return str(self._get_property(device, 'DriveConnectionInterface'))
        
    def is_systeminternal(self, device):
        return bool(self._get_property(device, 'DeviceIsSystemInternal'))

    def is_handleable(self, device):
        """Should this device be handled by udisks?

        Currently this just means that the device is removable and holds a
        filesystem."""

        if self.is_filesystem(device) and not self.is_systeminternal(device):
            return True
        else:
            return False
    # --- End PARAMETERS ---

    #--- ACTIONS ---
    def mount(self, device, options):
        if not self.is_mounted(device) and self.is_handleable(device):
            filesystem = self._get_property(device, 'IdType')
            dev_obj = self.bus.get_object(UDISKS_OBJECT, device)
            dev = dbus.Interface(dev_obj, DBUS_PROPS_INTERFACE)
            try:
                dev.FilesystemMount(filesystem, options,
                        dbus_interface=UDISKS_DEVICE_INTERFACE)
            except dbus.DBusException as e:
                return (1, e)

            return (0, '')
        else:
            return (1, 'Device %s is already mount or not handleable' % self.devname(device))


    def unmount(self, device):
        if self.is_mounted(device) and self.is_handleable(device):
            dev_obj = self.bus.get_object(UDISKS_OBJECT, device)
            dev = dbus.Interface(dev_obj, DBUS_PROPS_INTERFACE)
            try: 
                dev.FilesystemUnmount([], dbus_interface=UDISKS_DEVICE_INTERFACE)
            except dbus.DBusException as e:
                return (1, e)

            return (0, '')
        else:
            return (1, 'Device %s is not mounted or not handleable' % self.devname(device))

    def get_mountable_devices(self):
        devlist = self.get_udisk_devs()
        
        d = []
        for i in devlist:
            dev_obj = self.bus.get_object(UDISKS_OBJECT, i)
            dev = dbus.Interface(dev_obj, DBUS_PROPS_INTERFACE)
            if not self.is_mounted(i) and self.is_handleable(i) and self.has_media(i):
                d.append(str(self.deviface(i)).upper() + ': ' +
                    str(self.devlabel(i)) + ' (' + str(self.devname(i)) + ')')
        return d

    def get_unmountable_devices(self):
        devlist = self.get_udisk_devs()
        
        d = []
        for i in devlist:
            dev_obj = self.bus.get_object(UDISKS_OBJECT, i)
            dev = dbus.Interface(dev_obj, DBUS_PROPS_INTERFACE)
            if self.is_mounted(i) and self.is_handleable(i):
                mountpoins = str(', '.join(self.mount_paths(i)))
                d.append('"' + mountpoins + '"' + ' (' + str(self.devname(i)) + ')')
        return d

    def nametodev(self, name):
        return str(self.ud_manager.FindDeviceByDeviceFile(name))
    #--- End ACTIONS ---

# --- FUNCTIONS ---
def send_Notification(title, message, icon='dialog-warning',
                    urgency=URGENCY_CRITICAL, timeout=5, category='device'):
    '''Notify:
    send_Notification title, message '''
    import pynotify

    pynotify.init(title)

    notify = pynotify.Notification(title, message, icon)
    notify.set_urgency(urgency)
    notify.set_timeout(timeout * 1000) # 5 seconds
    notify.set_category(category)

    res = notify.show()
    pynotify.uninit()

    return res

def _run_Cmd(cmd, string):
    from subprocess import Popen
    from subprocess import PIPE

    returncode = 1
    
    try:
        p = Popen(cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE, close_fds=True)
        p.stdin.writelines(string)
    
        (out, err) = p.communicate()
        returncode = p.returncode
    except:
        (out, err) = ('', 'Error while exec: %s' % str(cmd))

    return (returncode, out, err)

def Mount():
    DMENU_MOUNT = ['dmenu', '-fn', 'DroidSans-9', '-sb', '#525252', '-p', 'Mount:']

    U = Udisks()
    Devs = U.get_mountable_devices()

    if Devs:
        string = '\n'.join(Devs)
        mydev = _run_Cmd(DMENU_MOUNT, string)
        if mydev[0] is 0:
            devname = re.search(ur"\([^}]+\)", mydev[1]).group().strip(')(')
            dev = U.nametodev(devname)
            if dev and dev in U.Properties.keys():
                (res, err) = U.mount(dev, [])
                if res is 0:
                    send_Notification('Mount', 'Mount "%s" succsessful!' % devname,
                        icon=ICON_SUCCESS, urgency=URGENCY_NORMAL)
                else:
                    send_Notification('Mount', 'Mount "%s" ERROR!\n%s' % (devname, err),
                        icon=ICON_ERROR, urgency=URGENCY_CRITICAL)
            else:
                raise IOError('Error mounting. Unknown device')
        else:
            print 'Error runing dmenu'

def Unmount():
    DMENU_UNMOUNT = ['dmenu', '-p', 'Umount:']

    U = Udisks()
    Devs = U.get_unmountable_devices()
    
    if Devs:
        string = '\n'.join(Devs)
        mydev = _run_Cmd(DMENU_UNMOUNT, string)
        if mydev[0] is 0:
            devname = re.search(ur"\([^}]+\)", mydev[1]).group().strip(')(')
            dev = U.nametodev(devname)
            if dev and dev in U.Properties.keys():
                (res, err) = U.unmount(dev)
                if res is 0:
                    send_Notification('Unmount', 'Unmount %s sucsessfull' % devname,
                        icon=ICON_SUCCESS, urgency=URGENCY_NORMAL)
                else:
                    send_Notification('Unmount', 'Unmount "%s" ERROR!\n%s' % (devname, err),
                        icon=ICON_ERROR, urgency=URGENCY_CRITICAL)
            else:
                raise IOError('Error unmounting. Unknown device')
        else:
            print 'Error runing dmenu'


# --- End FUNCTIONS ---

# ----- MAIN ----
def main():
    import optparse
    parser = optparse.OptionParser()
    parser.add_option('--mount', action='store_true',
                      dest='mount', default=False,
                      help='Mount devices')
    parser.add_option('--unmount', action='store_true',
                      dest='unmount', default=False,
                      help='Unmount devices')
    (options, args) = parser.parse_args()

    if options.unmount:
        Unmount()
    if options.mount:
        res = Mount()
# ----- End MAIN ----

if __name__ == '__main__':
    main()
